# -*- coding: utf-8 -*-
"""
-------------------------------------------------
File Name： normalize
Description :
Author : 'li'
date： 2022/1/28
-------------------------------------------------
Change Activity:
2022/1/28:
-------------------------------------------------
"""
import numpy as np
from torch import Tensor


class Normalize:
    def __init__(self, mean=(127, 127, 127), std=(40, 40, 40), to_float32=True, input_key='img', output_key='img'):
        """

        Args:
            mean:
            std: standard deviation
            to_float32:
        """
        self.mean = mean
        self.std = std
        self.to_float32 = to_float32
        self.input_key = input_key
        self.output_key = output_key

    def __call__(self, data):
        if isinstance(data, Tensor):
            raise Exception('not implement')
        else:
            img = data[self.input_key]
            if self.to_float32:
                img = img.astype(np.float32)
            mean = np.array(self.mean)
            std = np.array(self.std)
            img = (img - mean) / std
            data[self.output_key] = img
            return data
